{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "89c9afa4-399c-4a64-ac27-dd1aeecbc8f3",
   "metadata": {},
   "source": [
    "# Fast-Fourier-Transformation Example\n",
    "\n",
    "## Machine-Learned Inverse FFT\n",
    "\n",
    "Let's start with a sinusoidal curve: that's the input (`x`) that we want to learn, given its FFT (`y`).\n",
    "The goal is to do an `InverseRealFFT` by gradient descent.\n",
    "\n",
    "For this problem the input (`x`) is real (`float32`) and label (`y`), the FFT, is complex (`complex64`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2a73f9ca-19a5-44e8-96c1-c488ecf356c0",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\t- Added replace rule for module \"github.com/gomlx/gomlx\" to local directory \"/home/janpf/Projects/gomlx\".\n"
     ]
    }
   ],
   "source": [
    "!*rm -f go.work && go work init && go work use . \"${HOME}/Projects/gomlx\"\n",
    "%goworkfix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e34bf9b0-a4ed-4afe-a52d-deccd7bd4620",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "x: shape=(Float32)[1 100]\n",
      "y: shape=(Complex64)[1 51]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<svg preserveAspectRatio=\"xMidYMid meet\" xmlns=\"http://www.w3.org/2000/svg\" width=\"1024\" height=\"320\" viewbox=\"0 0 1024 320\" style=\"background-color:#f8f8f8\"><defs><marker refX=\"5\" refY=\"5\" markerUnits=\"userSpaceOnUse\" markerWidth=\"2%\" markerHeight=\"2%\" id=\"circle\" viewBox=\"0 0 10 10 \"><circle cx=\"5\" cy=\"5\" r=\"3\" fill=\"none\" stroke=\"black\"/></marker><marker markerWidth=\"2%\" markerHeight=\"2%\" id=\"filled-circle\" viewBox=\"0 0 10 10 \" refX=\"5\" refY=\"5\" markerUnits=\"userSpaceOnUse\"><circle cx=\"5\" cy=\"5\" r=\"3\" fill=\"black\" stroke=\"none\"/></marker><marker id=\"square\" viewBox=\"0 0 10 10 \" refX=\"5\" refY=\"5\" markerUnits=\"userSpaceOnUse\" markerWidth=\"2%\" markerHeight=\"2%\"><rect height=\"6\" fill=\"none\" stroke=\"black\" x=\"2\" y=\"2\" width=\"6\"/></marker><marker markerHeight=\"2%\" id=\"filled-square\" viewBox=\"0 0 10 10 \" refX=\"5\" refY=\"5\" markerUnits=\"userSpaceOnUse\" markerWidth=\"2%\"><rect fill=\"black\" stroke=\"none\" x=\"2\" y=\"2\" width=\"6\" height=\"6\"/></marker></defs><g transform=\"translate(70 250 )scale(1 -1 )\" fill=\"none\" marker-end=\"url(#square)\" stroke-width=\"2px\" marker-mid=\"url(#square)\" stroke-linecap=\"round\" stroke-linejoin=\"round\" marker-start=\"url(#square)\" stroke=\"hsl(90, 47%, 65%)\"><path vector-effect=\"non-scaling-stroke\" d=\"M1.768000e+01,9.000000e+01 L2.625212e+01,1.008502e+02 L3.482424e+01,1.115293e+02 L4.339636e+01,1.218689e+02 L5.196848e+01,1.317058e+02 L6.054061e+01,1.408851e+02 L6.911273e+01,1.492618e+02 L7.768485e+01,1.567040e+02 L8.625697e+01,1.630942e+02 L9.482909e+01,1.683316e+02 L1.034012e+02,1.723338e+02 L1.119733e+02,1.750374e+02 L1.205455e+02,1.764000e+02 L1.291176e+02,1.764000e+02 L1.376897e+02,1.750374e+02 L1.462618e+02,1.723337e+02 L1.548339e+02,1.683316e+02 L1.634061e+02,1.630942e+02 L1.719782e+02,1.567040e+02 L1.805503e+02,1.492618e+02 L1.891224e+02,1.408851e+02 L1.976945e+02,1.317058e+02 L2.062667e+02,1.218688e+02 L2.148388e+02,1.115293e+02 L2.234109e+02,1.008502e+02 L2.319830e+02,9.000000e+01 L2.405552e+02,7.914979e+01 L2.491273e+02,6.847070e+01 L2.576994e+02,5.813114e+01 L2.662715e+02,4.829417e+01 L2.748436e+02,3.911494e+01 L2.834158e+02,3.073819e+01 L2.919879e+02,2.329603e+01 L3.005600e+02,1.690582e+01 L3.091321e+02,1.166837e+01 L3.177042e+02,7.666243e+00 L3.262764e+02,4.962576e+00 L3.348485e+02,3.600005e+00 L3.434206e+02,3.600005e+00 L3.519927e+02,4.962587e+00 L3.605648e+02,7.666258e+00 L3.691370e+02,1.166838e+01 L3.777091e+02,1.690585e+01 L3.862812e+02,2.329605e+01 L3.948533e+02,3.073821e+01 L4.034255e+02,3.911497e+01 L4.119976e+02,4.829422e+01 L4.205697e+02,5.813120e+01 L4.291418e+02,6.847076e+01 L4.377139e+02,7.914981e+01 L4.462861e+02,9.000002e+01 L4.548582e+02,1.008502e+02 L4.634303e+02,1.115293e+02 L4.720024e+02,1.218689e+02 L4.805745e+02,1.317058e+02 L4.891467e+02,1.408851e+02 L4.977188e+02,1.492618e+02 L5.062909e+02,1.567040e+02 L5.148630e+02,1.630942e+02 L5.234352e+02,1.683316e+02 L5.320073e+02,1.723338e+02 L5.405794e+02,1.750374e+02 L5.491515e+02,1.764000e+02 L5.577236e+02,1.764000e+02 L5.662958e+02,1.750374e+02 L5.748679e+02,1.723337e+02 L5.834400e+02,1.683316e+02 L5.920121e+02,1.630941e+02 L6.005842e+02,1.567039e+02 L6.091564e+02,1.492618e+02 L6.177285e+02,1.408850e+02 L6.263006e+02,1.317058e+02 L6.348727e+02,1.218688e+02 L6.434448e+02,1.115292e+02 L6.520170e+02,1.008501e+02 L6.605891e+02,9.000000e+01 L6.691612e+02,7.914979e+01 L6.777333e+02,6.847070e+01 L6.863055e+02,5.813114e+01 L6.948776e+02,4.829417e+01 L7.034497e+02,3.911493e+01 L7.120218e+02,3.073817e+01 L7.205939e+02,2.329602e+01 L7.291661e+02,1.690582e+01 L7.377382e+02,1.166836e+01 L7.463103e+02,7.666243e+00 L7.548824e+02,4.962576e+00 L7.634545e+02,3.600000e+00 L7.720267e+02,3.600010e+00 L7.805988e+02,4.962597e+00 L7.891709e+02,7.666274e+00 L7.977430e+02,1.166841e+01 L8.063152e+02,1.690588e+01 L8.148873e+02,2.329609e+01 L8.234594e+02,3.073825e+01 L8.320315e+02,3.911502e+01 L8.406036e+02,4.829427e+01 L8.491758e+02,5.813117e+01 L8.577479e+02,6.847074e+01 L8.663200e+02,7.914983e+01 \"/></g><g fill=\"none\" stroke-linecap=\"round\" stroke-linejoin=\"round\" stroke=\"black\" stroke-width=\"2px\" transform=\"translate(70 250 )scale(1 -1 )\"><path vector-effect=\"non-scaling-stroke\" d=\"M1.768000e+01,0 L1.768000e+01,-6 M1.034012e+02,0 L1.034012e+02,-6 M1.891224e+02,0 L1.891224e+02,-6 M2.748436e+02,0 L2.748436e+02,-6 M3.605648e+02,0 L3.605648e+02,-6 M4.462861e+02,0 L4.462861e+02,-6 M5.320073e+02,0 L5.320073e+02,-6 M6.177285e+02,0 L6.177285e+02,-6 M7.034497e+02,0 L7.034497e+02,-6 M7.891709e+02,0 L7.891709e+02,-6 \"/></g><g font-family=\"sans-serif\" font-weight=\"normal\" text-anchor=\"middle\" stroke-linecap=\"round\" stroke-linejoin=\"round\" font-style=\"normal\" dominant-baseline=\"hanging\" stroke=\"black\" fill=\"black\" font-size=\"12px\" stroke-width=\"2px\" transform=\"translate(70 250 )scale(1 1 )\"><text stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"1.768000e+01\" y=\"10\" dominant-baseline=\"hanging\">0</text><text x=\"1.034012e+02\" y=\"10\" dominant-baseline=\"hanging\" stroke=\"none\" vector-effect=\"non-scaling-stroke\">10</text><text y=\"10\" dominant-baseline=\"hanging\" stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"1.891224e+02\">20</text><text vector-effect=\"non-scaling-stroke\" x=\"2.748436e+02\" y=\"10\" dominant-baseline=\"hanging\" stroke=\"none\">30</text><text stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"3.605648e+02\" y=\"10\" dominant-baseline=\"hanging\">40</text><text x=\"4.462861e+02\" y=\"10\" dominant-baseline=\"hanging\" stroke=\"none\" vector-effect=\"non-scaling-stroke\">50</text><text y=\"10\" dominant-baseline=\"hanging\" stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"5.320073e+02\">60</text><text vector-effect=\"non-scaling-stroke\" x=\"6.177285e+02\" y=\"10\" dominant-baseline=\"hanging\" stroke=\"none\">70</text><text vector-effect=\"non-scaling-stroke\" x=\"7.034497e+02\" y=\"10\" dominant-baseline=\"hanging\" stroke=\"none\">80</text><text stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"7.891709e+02\" y=\"10\" dominant-baseline=\"hanging\">90</text></g><g stroke=\"black\" stroke-width=\"2px\" font-family=\"sans-serif\" font-weight=\"bold\" stroke-linecap=\"round\" dominant-baseline=\"baseline\" transform=\"translate(70 250 )scale(1 1 )rotate(0 0 0 )\" text-anchor=\"middle\" fill=\"black\" stroke-linejoin=\"round\" font-size=\"12px\" font-style=\"normal\"><text y=\"-6\" dominant-baseline=\"baseline\" stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"442\">Steps</text></g><g stroke-linejoin=\"round\" font-style=\"normal\" stroke=\"black\" stroke-width=\"2px\" transform=\"translate(70 250 )scale(1 -1 )\" font-family=\"sans-serif\" font-weight=\"bold\" text-anchor=\"middle\" font-size=\"12px\" dominant-baseline=\"baseline\" fill=\"black\" stroke-linecap=\"round\"><path vector-effect=\"non-scaling-stroke\" d=\"M0,9.000000e+01 L-6,9.000000e+01 \"/></g><g fill=\"black\" stroke-linecap=\"round\" stroke-linejoin=\"round\" font-family=\"sans-serif\" font-size=\"12px\" font-style=\"normal\" dominant-baseline=\"middle\" stroke=\"black\" stroke-width=\"2px\" transform=\"translate(70 250 )scale(1 1 )\" font-weight=\"normal\" text-anchor=\"end\"><text vector-effect=\"non-scaling-stroke\" x=\"-10\" y=\"-9.000000e+01\" dominant-baseline=\"middle\" stroke=\"none\">-0.000</text></g><g stroke-width=\"0.5px\" transform=\"translate(70 250 )scale(1 -1 )\" font-family=\"sans-serif\" font-weight=\"normal\" fill=\"black\" stroke-linejoin=\"round\" dominant-baseline=\"middle\" text-anchor=\"end\" stroke-linecap=\"round\" font-size=\"12px\" font-style=\"normal\" stroke=\"gray\"><path d=\"M0,9.000000e+01 L884,9.000000e+01 \" vector-effect=\"non-scaling-stroke\"/></g><g font-size=\"12px\" stroke=\"black\" stroke-width=\"2px\" font-family=\"sans-serif\" text-anchor=\"end\" fill=\"none\" stroke-linejoin=\"round\" font-style=\"normal\" dominant-baseline=\"middle\" font-weight=\"normal\" stroke-linecap=\"round\"><rect vector-effect=\"non-scaling-stroke\" x=\"70\" y=\"70\" width=\"884\" height=\"180\"/></g></svg>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import (\n",
    "    \"github.com/gomlx/gopjrt/dtypes\"\n",
    "    . \"github.com/gomlx/gomlx/pkg/core/graph\"\n",
    "    . \"github.com/gomlx/exceptions\"\n",
    "    \"github.com/gomlx/gomlx/pkg/core/tensors\"\n",
    "    \"github.com/gomlx/gomlx/pkg/support/xslices\"\n",
    "    mg \"github.com/gomlx/gomlx/ui/gonb/margaid\"\n",
    "    \"github.com/janpfeifer/gonb/gonbui\"\n",
    "\n",
    "    _ \"github.com/gomlx/gomlx/backends/default\"\n",
    ")\n",
    "\n",
    "// manager always created at initialization.\n",
    "var backend = backends.MustNew()\n",
    "\n",
    "const (\n",
    "    NumPoints = 100\n",
    "    Frequency = 2.0  // Number of curves in samples.\n",
    "    RealDType = dtypes.Float32\n",
    "    ComplexDType = dtypes.Complex64\n",
    ")\n",
    "\n",
    "// CalculateXY returns (x, y) of our problem, where y is a sinusoidal curve and x is its FFT.\n",
    "func CalculateXY() (x, y *tensors.Tensor) {\n",
    "    e := MustNewExec(backend, func (g *Graph) (x, y *Node) {\n",
    "        x = Iota(g, shapes.Make(RealDType, 1, NumPoints), 1)\n",
    "        x = MulScalar(x, 2.0*math.Pi*Frequency/float64(NumPoints))\n",
    "        x = Sin(x)\n",
    "        y = RealFFT(x)\n",
    "        return\n",
    "    })\n",
    "    return e.MustExec2()\n",
    "}\n",
    "\n",
    "func Plot(displayId string, width, height int, coordinates []*tensors.Tensor, names []string) {\n",
    "    plts := mg.New(width, height)\n",
    "    for ii, t := range coordinates {\n",
    "        var values []float64\n",
    "        switch t.DType() {\n",
    "        case dtypes.Float64:\n",
    "            values = tensors.CopyFlatData[float64](t)\n",
    "        case dtypes.Float32:\n",
    "            values32 := tensors.CopyFlatData[float32](t)\n",
    "            values = xslices.Map(values32, func (v float32) float64 { return float64(v) })\n",
    "        default:\n",
    "            Panicf(\"only float32 and float64 tensor dtypes are accepted by Plot, got t.shape=%s\", t.Shape())\n",
    "        }\n",
    "        var name string\n",
    "        if len(names) > ii {\n",
    "            name = names[ii]\n",
    "        }\n",
    "        plts.AddValues(name, \"\", values)\n",
    "    }\n",
    "    if displayId == \"\" {\n",
    "        plts.Plot()\n",
    "    } else {\n",
    "        gonbui.UpdateHTML(displayId, plts.PlotToHTML())\n",
    "    }\n",
    "}\n",
    "\n",
    "%%\n",
    "x, y := CalculateXY()\n",
    "fmt.Printf(\"x: shape=%s\\n\", x.Shape())\n",
    "fmt.Printf(\"y: shape=%s\\n\", y.Shape())\n",
    "Plot(\"\", 1024, 320, []*tensors.Tensor{x}, nil)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "749d6fbd-c7ad-4bec-80d4-8a2dd1dd1531",
   "metadata": {},
   "source": [
    "### Train the model\n",
    "\n",
    "If you run it, you'll see the plot of the \"learnedX\" adjusting towards \"x\", the original sinusoidal curve."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "555b0279-e565-4857-a422-6ed214305408",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<svg height=\"356\" viewbox=\"0 0 1024 320\" style=\"background-color:#f8f8f8\" preserveAspectRatio=\"xMidYMid meet\" xmlns=\"http://www.w3.org/2000/svg\" width=\"1024\"><defs><marker refX=\"5\" refY=\"5\" markerUnits=\"userSpaceOnUse\" markerWidth=\"2%\" markerHeight=\"2%\" id=\"circle\" viewBox=\"0 0 10 10 \"><circle cy=\"5\" r=\"3\" fill=\"none\" stroke=\"black\" cx=\"5\"/></marker><marker markerUnits=\"userSpaceOnUse\" markerWidth=\"2%\" markerHeight=\"2%\" id=\"filled-circle\" viewBox=\"0 0 10 10 \" refX=\"5\" refY=\"5\"><circle cx=\"5\" cy=\"5\" r=\"3\" fill=\"black\" stroke=\"none\"/></marker><marker refY=\"5\" markerUnits=\"userSpaceOnUse\" markerWidth=\"2%\" markerHeight=\"2%\" id=\"square\" viewBox=\"0 0 10 10 \" refX=\"5\"><rect x=\"2\" y=\"2\" width=\"6\" height=\"6\" fill=\"none\" stroke=\"black\"/></marker><marker refY=\"5\" markerUnits=\"userSpaceOnUse\" markerWidth=\"2%\" markerHeight=\"2%\" id=\"filled-square\" viewBox=\"0 0 10 10 \" refX=\"5\"><rect x=\"2\" y=\"2\" width=\"6\" height=\"6\" fill=\"black\" stroke=\"none\"/></marker></defs><g stroke-width=\"2px\" stroke-linecap=\"round\" marker-mid=\"url(#square)\" stroke=\"hsl(90, 47%, 65%)\" fill=\"none\" marker-end=\"url(#square)\" transform=\"translate(70 250 )scale(1 -1 )\" stroke-linejoin=\"round\" marker-start=\"url(#square)\"><path vector-effect=\"non-scaling-stroke\" d=\"M1.768000e+01,9.018442e+01 L2.625212e+01,1.009365e+02 L3.482424e+01,1.116592e+02 L4.339636e+01,1.218885e+02 L5.196848e+01,1.317009e+02 L6.054061e+01,1.406720e+02 L6.911273e+01,1.492110e+02 L7.768485e+01,1.566179e+02 L8.625697e+01,1.630492e+02 L9.482909e+01,1.681726e+02 L1.034012e+02,1.722598e+02 L1.119733e+02,1.748783e+02 L1.205455e+02,1.764000e+02 L1.291176e+02,1.761877e+02 L1.376897e+02,1.748893e+02 L1.462618e+02,1.722069e+02 L1.548339e+02,1.681859e+02 L1.634061e+02,1.630325e+02 L1.719782e+02,1.566730e+02 L1.805503e+02,1.491776e+02 L1.891224e+02,1.408337e+02 L1.976945e+02,1.316888e+02 L2.062667e+02,1.218477e+02 L2.148388e+02,1.114828e+02 L2.234109e+02,1.008871e+02 L2.319830e+02,9.001868e+01 L2.405552e+02,7.924377e+01 L2.491273e+02,6.854082e+01 L2.576994e+02,5.829190e+01 L2.662715e+02,4.848711e+01 L2.748436e+02,3.921227e+01 L2.834158e+02,3.083985e+01 L2.919879e+02,2.350122e+01 L3.005600e+02,1.702543e+01 L3.091321e+02,1.182844e+01 L3.177042e+02,7.997745e+00 L3.262764e+02,5.378169e+00 L3.348485e+02,3.842576e+00 L3.434206e+02,3.821945e+00 L3.519927e+02,5.160807e+00 L3.605648e+02,7.844087e+00 L3.691370e+02,1.189244e+01 L3.777091e+02,1.690206e+01 L3.862812e+02,2.340879e+01 L3.948533e+02,3.106863e+01 L4.034255e+02,3.937168e+01 L4.119976e+02,4.837920e+01 L4.205697e+02,5.842303e+01 L4.291418e+02,6.861834e+01 L4.377139e+02,7.899020e+01 L4.462861e+02,9.007014e+01 L4.548582e+02,1.008416e+02 L4.634303e+02,1.114966e+02 L4.720024e+02,1.218146e+02 L4.805745e+02,1.315460e+02 L4.891467e+02,1.408036e+02 L4.977188e+02,1.493727e+02 L5.062909e+02,1.566447e+02 L5.148630e+02,1.630798e+02 L5.234352e+02,1.682840e+02 L5.320073e+02,1.719686e+02 L5.405794e+02,1.748051e+02 L5.491515e+02,1.763590e+02 L5.577236e+02,1.763492e+02 L5.662958e+02,1.750504e+02 L5.748679e+02,1.721647e+02 L5.834400e+02,1.680819e+02 L5.920121e+02,1.628052e+02 L6.005842e+02,1.567074e+02 L6.091564e+02,1.488831e+02 L6.177285e+02,1.407857e+02 L6.263006e+02,1.318675e+02 L6.348727e+02,1.218919e+02 L6.434448e+02,1.114634e+02 L6.520170e+02,1.009062e+02 L6.605891e+02,8.992699e+01 L6.691612e+02,7.926046e+01 L6.777333e+02,6.861514e+01 L6.863055e+02,5.802950e+01 L6.948776e+02,4.830800e+01 L7.034497e+02,3.944058e+01 L7.120218e+02,3.072777e+01 L7.205939e+02,2.353566e+01 L7.291661e+02,1.703110e+01 L7.377382e+02,1.176467e+01 L7.463103e+02,7.937511e+00 L7.548824e+02,5.170695e+00 L7.634545e+02,3.600000e+00 L7.720267e+02,3.897603e+00 L7.805988e+02,5.187952e+00 L7.891709e+02,7.939000e+00 L7.977430e+02,1.177236e+01 L8.063152e+02,1.710991e+01 L8.148873e+02,2.354074e+01 L8.234594e+02,3.091920e+01 L8.320315e+02,3.935715e+01 L8.406036e+02,4.850126e+01 L8.491758e+02,5.813852e+01 L8.577479e+02,6.857974e+01 L8.663200e+02,7.916387e+01 \"/></g><g stroke-width=\"2px\" marker-end=\"url(#square)\" transform=\"translate(70 250 )scale(1 -1 )\" stroke=\"hsl(301, 88%, 65%)\" marker-start=\"url(#square)\" fill=\"none\" stroke-linecap=\"round\" marker-mid=\"url(#square)\" stroke-linejoin=\"round\"><path vector-effect=\"non-scaling-stroke\" d=\"M1.768000e+01,9.004901e+01 L2.625212e+01,1.008779e+02 L3.482424e+01,1.115361e+02 L4.339636e+01,1.218554e+02 L5.196848e+01,1.316730e+02 L6.054061e+01,1.408343e+02 L6.911273e+01,1.491946e+02 L7.768485e+01,1.566222e+02 L8.625697e+01,1.629998e+02 L9.482909e+01,1.682270e+02 L1.034012e+02,1.722213e+02 L1.119733e+02,1.749197e+02 L1.205455e+02,1.762796e+02 L1.291176e+02,1.762796e+02 L1.376897e+02,1.749197e+02 L1.462618e+02,1.722213e+02 L1.548339e+02,1.682270e+02 L1.634061e+02,1.629998e+02 L1.719782e+02,1.566222e+02 L1.805503e+02,1.491946e+02 L1.891224e+02,1.408343e+02 L1.976945e+02,1.316730e+02 L2.062667e+02,1.218553e+02 L2.148388e+02,1.115361e+02 L2.234109e+02,1.008779e+02 L2.319830e+02,9.004900e+01 L2.405552e+02,7.922008e+01 L2.491273e+02,6.856193e+01 L2.576994e+02,5.824265e+01 L2.662715e+02,4.842497e+01 L2.748436e+02,3.926375e+01 L2.834158e+02,3.090342e+01 L2.919879e+02,2.347585e+01 L3.005600e+02,1.709818e+01 L3.091321e+02,1.187100e+01 L3.177042e+02,7.876722e+00 L3.262764e+02,5.178358e+00 L3.348485e+02,3.818459e+00 L3.434206e+02,3.818459e+00 L3.519927e+02,5.178368e+00 L3.605648e+02,7.876737e+00 L3.691370e+02,1.187101e+01 L3.777091e+02,1.709821e+01 L3.862812e+02,2.347588e+01 L3.948533e+02,3.090344e+01 L4.034255e+02,3.926378e+01 L4.119976e+02,4.842502e+01 L4.205697e+02,5.824270e+01 L4.291418e+02,6.856199e+01 L4.377139e+02,7.922010e+01 L4.462861e+02,9.004903e+01 L4.548582e+02,1.008780e+02 L4.634303e+02,1.115361e+02 L4.720024e+02,1.218554e+02 L4.805745e+02,1.316731e+02 L4.891467e+02,1.408343e+02 L4.977188e+02,1.491946e+02 L5.062909e+02,1.566222e+02 L5.148630e+02,1.629999e+02 L5.234352e+02,1.682270e+02 L5.320073e+02,1.722213e+02 L5.405794e+02,1.749197e+02 L5.491515e+02,1.762796e+02 L5.577236e+02,1.762796e+02 L5.662958e+02,1.749196e+02 L5.748679e+02,1.722213e+02 L5.834400e+02,1.682270e+02 L5.920121e+02,1.629998e+02 L6.005842e+02,1.566221e+02 L6.091564e+02,1.491946e+02 L6.177285e+02,1.408342e+02 L6.263006e+02,1.316730e+02 L6.348727e+02,1.218553e+02 L6.434448e+02,1.115360e+02 L6.520170e+02,1.008779e+02 L6.605891e+02,9.004901e+01 L6.691612e+02,7.922008e+01 L6.777333e+02,6.856193e+01 L6.863055e+02,5.824265e+01 L6.948776e+02,4.842497e+01 L7.034497e+02,3.926373e+01 L7.120218e+02,3.090340e+01 L7.205939e+02,2.347584e+01 L7.291661e+02,1.709818e+01 L7.377382e+02,1.187099e+01 L7.463103e+02,7.876722e+00 L7.548824e+02,5.178358e+00 L7.634545e+02,3.818454e+00 L7.720267e+02,3.818464e+00 L7.805988e+02,5.178379e+00 L7.891709e+02,7.876753e+00 L7.977430e+02,1.187104e+01 L8.063152e+02,1.709824e+01 L8.148873e+02,2.347591e+01 L8.234594e+02,3.090349e+01 L8.320315e+02,3.926382e+01 L8.406036e+02,4.842508e+01 L8.491758e+02,5.824268e+01 L8.577479e+02,6.856197e+01 L8.663200e+02,7.922011e+01 \"/></g><g stroke-linecap=\"round\" transform=\"translate(70 250 )scale(1 -1 )\" stroke=\"black\" stroke-linejoin=\"round\" fill=\"none\" stroke-width=\"2px\"><path vector-effect=\"non-scaling-stroke\" d=\"M1.768000e+01,0 L1.768000e+01,-6 M1.034012e+02,0 L1.034012e+02,-6 M1.891224e+02,0 L1.891224e+02,-6 M2.748436e+02,0 L2.748436e+02,-6 M3.605648e+02,0 L3.605648e+02,-6 M4.462861e+02,0 L4.462861e+02,-6 M5.320073e+02,0 L5.320073e+02,-6 M6.177285e+02,0 L6.177285e+02,-6 M7.034497e+02,0 L7.034497e+02,-6 M7.891709e+02,0 L7.891709e+02,-6 \"/></g><g stroke-linejoin=\"round\" font-size=\"12px\" font-style=\"normal\" stroke-width=\"2px\" stroke-linecap=\"round\" transform=\"translate(70 250 )scale(1 1 )\" font-family=\"sans-serif\" font-weight=\"normal\" dominant-baseline=\"hanging\" text-anchor=\"middle\" fill=\"black\" stroke=\"black\"><text y=\"10\" dominant-baseline=\"hanging\" stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"1.768000e+01\">0</text><text stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"1.034012e+02\" y=\"10\" dominant-baseline=\"hanging\">10</text><text y=\"10\" dominant-baseline=\"hanging\" stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"1.891224e+02\">20</text><text x=\"2.748436e+02\" y=\"10\" dominant-baseline=\"hanging\" stroke=\"none\" vector-effect=\"non-scaling-stroke\">30</text><text dominant-baseline=\"hanging\" stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"3.605648e+02\" y=\"10\">40</text><text vector-effect=\"non-scaling-stroke\" x=\"4.462861e+02\" y=\"10\" dominant-baseline=\"hanging\" stroke=\"none\">50</text><text stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"5.320073e+02\" y=\"10\" dominant-baseline=\"hanging\">60</text><text stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"6.177285e+02\" y=\"10\" dominant-baseline=\"hanging\">70</text><text stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"7.034497e+02\" y=\"10\" dominant-baseline=\"hanging\">80</text><text x=\"7.891709e+02\" y=\"10\" dominant-baseline=\"hanging\" stroke=\"none\" vector-effect=\"non-scaling-stroke\">90</text></g><g stroke-linecap=\"round\" transform=\"translate(70 250 )scale(1 1 )rotate(0 0 0 )\" font-family=\"sans-serif\" dominant-baseline=\"baseline\" font-style=\"normal\" fill=\"black\" stroke-width=\"2px\" font-weight=\"bold\" stroke=\"black\" stroke-linejoin=\"round\" font-size=\"12px\" text-anchor=\"middle\"><text y=\"-6\" dominant-baseline=\"baseline\" stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"442\">Steps</text></g><g stroke-linecap=\"round\" font-family=\"sans-serif\" font-weight=\"bold\" dominant-baseline=\"baseline\" transform=\"translate(70 250 )scale(1 -1 )\" stroke=\"black\" stroke-linejoin=\"round\" font-size=\"12px\" font-style=\"normal\" text-anchor=\"middle\" fill=\"black\" stroke-width=\"2px\"><path d=\"M0,3.647966e+00 L-6,3.647966e+00 M0,9.004901e+01 L-6,9.004901e+01 \" vector-effect=\"non-scaling-stroke\"/></g><g stroke=\"black\" stroke-linejoin=\"round\" font-style=\"normal\" text-anchor=\"end\" fill=\"black\" stroke-linecap=\"round\" transform=\"translate(70 250 )scale(1 1 )\" dominant-baseline=\"middle\" font-size=\"12px\" stroke-width=\"2px\" font-family=\"sans-serif\" font-weight=\"normal\"><text vector-effect=\"non-scaling-stroke\" x=\"-10\" y=\"-3.647966e+00\" dominant-baseline=\"middle\" stroke=\"none\">-1.000</text><text vector-effect=\"non-scaling-stroke\" x=\"-10\" y=\"-9.004901e+01\" dominant-baseline=\"middle\" stroke=\"none\">0.000</text></g><g transform=\"translate(70 250 )scale(1 -1 )\" font-weight=\"normal\" stroke-linejoin=\"round\" font-size=\"12px\" font-family=\"sans-serif\" dominant-baseline=\"middle\" stroke=\"gray\" font-style=\"normal\" text-anchor=\"end\" fill=\"black\" stroke-width=\"0.5px\" stroke-linecap=\"round\"><path d=\"M0,3.647966e+00 L884,3.647966e+00 M0,9.004901e+01 L884,9.004901e+01 \" vector-effect=\"non-scaling-stroke\"/></g><g font-style=\"normal\" text-anchor=\"end\" stroke-linecap=\"round\" font-family=\"sans-serif\" font-weight=\"normal\" dominant-baseline=\"middle\" font-size=\"12px\" fill=\"none\" stroke-width=\"2px\" stroke=\"black\" stroke-linejoin=\"round\"><rect y=\"70\" width=\"884\" height=\"180\" vector-effect=\"non-scaling-stroke\" x=\"70\"/><g stroke-width=\"1px\" dominant-baseline=\"hanging\" stroke=\"hsl(90, 47%, 65%)\" text-anchor=\"start\" fill=\"hsl(90, 47%, 65%)\"><rect vector-effect=\"non-scaling-stroke\" x=\"76\" y=\"286\" width=\"12\" height=\"12\"/><g font-style=\"normal\" fill=\"black\" stroke-linecap=\"round\" font-family=\"sans-serif\" font-weight=\"normal\" stroke=\"black\" stroke-linejoin=\"round\" font-size=\"12px\"><text stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"92\" y=\"286\" dominant-baseline=\"hanging\">Learned</text><g stroke=\"hsl(301, 88%, 65%)\" text-anchor=\"start\" fill=\"hsl(301, 88%, 65%)\" stroke-width=\"1px\" dominant-baseline=\"hanging\"><rect x=\"76\" y=\"304\" width=\"12\" height=\"12\" vector-effect=\"non-scaling-stroke\"/><g font-style=\"normal\" fill=\"black\" stroke-linecap=\"round\" font-family=\"sans-serif\" font-weight=\"normal\" stroke=\"black\" stroke-linejoin=\"round\" font-size=\"12px\"><text stroke=\"none\" vector-effect=\"non-scaling-stroke\" x=\"92\" y=\"304\" dominant-baseline=\"hanging\">Truth</text></g></g></g></g></g></svg>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "      \u001b[1m 100% [========================================] (5095 steps/s)\u001b[0m [step=799] [loss+=0.0109] [~loss+=0.0666] [~loss=0.0666]                 \n"
     ]
    }
   ],
   "source": [
    "import (\n",
    "    . \"github.com/gomlx/gomlx/pkg/core/graph\"\n",
    "    \"github.com/gomlx/gomlx/pkg/ml/context\"\n",
    "    \"github.com/gomlx/gomlx/pkg/ml/datasets\"\n",
    "    \"github.com/gomlx/gomlx/pkg/ml/train\"\n",
    ")\n",
    "\n",
    "var (\n",
    "\tflagNumSteps     = flag.Int(\"steps\", 1000, \"Number of gradient descent steps to perform\")\n",
    "\tflagLearningRate = flag.Float64(\"learning_rate\", 0.1, \"Initial learning rate.\")\n",
    ")\n",
    "\n",
    "func TrainInverseRealFFT() {\n",
    "    x, y := CalculateXY()\n",
    "    ctx := context.New()\n",
    "\tctx.SetParam(optimizers.LearningRateKey, *flagLearningRate)\n",
    "\n",
    "    modelFn := func(ctx *context.Context, spec any, inputs []*Node) []*Node {\n",
    "        g := inputs[0].Graph()\n",
    "        learnedXVar := ctx.VariableWithShape(\"learnedX\", x.Shape())\n",
    "        predictedY := RealFFT(learnedXVar.ValueGraph(g))\n",
    "        return []*Node{predictedY}\n",
    "    }\n",
    "\n",
    "    dataset, err := data.InMemoryFromData(backend, \"dataset\", []any{x}, []any{y})\n",
    "    if err != nil {\n",
    "        panic(err)\n",
    "    }\n",
    "    dataset.BatchSize(1, false).Infinite(true)\n",
    "\n",
    "    opt := optimizers.Adam().Done()\n",
    "    trainer := train.NewTrainer(\n",
    "        backend, ctx, modelFn,\n",
    "        losses.MeanAbsoluteError,\n",
    "        opt,\n",
    "        nil, nil) // trainMetrics, evalMetrics\n",
    "\n",
    "\tloop := train.NewLoop(trainer)\n",
    "\tcommandline.AttachProgressBar(loop) // Attaches a progress bar to the loop.\n",
    "\n",
    "    // Plot learnedX\n",
    "    displayId := gonbui.UniqueID()\n",
    "    gonbui.UpdateHTML(displayId, \"\")\n",
    "    train.EveryNSteps(loop, 10, \"plot\", 0, func(loop *train.Loop, metrics []*tensors.Tensor) error {\n",
    "        learnedXVar := ctx.InspectVariable(context.RootScope, \"learnedX\")\n",
    "        learnedX := learnedXVar.Value()\n",
    "        Plot(displayId, 1024, 320, []*tensors.Tensor{x, learnedX}, []string{\"Truth\", \"Learned\"})\n",
    "        return nil\n",
    "    })\n",
    "    \n",
    "\t// Loop for given number of steps.\n",
    "\t_, err = loop.RunSteps(dataset, *flagNumSteps)\n",
    "\tif err != nil {\n",
    "\t\tpanic(err)\n",
    "\t}\n",
    "\n",
    "}\n",
    "\n",
    "%% --steps=800 --learning_rate=0.01\n",
    "TrainInverseRealFFT()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Go (gonb)",
   "language": "go",
   "name": "gonb"
  },
  "language_info": {
   "codemirror_mode": "",
   "file_extension": ".go",
   "mimetype": "text/x-go",
   "name": "go",
   "nbconvert_exporter": "",
   "pygments_lexer": "",
   "version": "go1.24.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
